/*
 * Copyright (c) 2024 endless-sky
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "firmware_upgrade_handler.h"
#include "utils_log.h"
#include "utils_string_ex.h"
#include "utils_crypto.h"
#include <chrono>
#include <iomanip>
#include <sstream>
#include <fstream>
#include <filesystem>
#include <openssl/md5.h>
#include "upgrade.h"

namespace El {
namespace Protocol {

void FirmwareUpgradeHandler::RegisterMethodHandlers(
    std::unordered_map<std::string, std::function<nlohmann::json(const std::string &, const nlohmann::json &)>>
        &handlers)
{
    // 注册固件升级相关方法
    handlers["firmware.upgrade"] = [this](const std::string &id, const nlohmann::json &params) {
        return HandleFirmwareUpgrade(id, params);
    };

    handlers["firmware.getStatus"] = [this](const std::string &id, const nlohmann::json &params) {
        return HandleGetUpgradeStatus(id, params);
    };

    handlers["firmware.upload"] = [this](const std::string &id, const nlohmann::json &params) {
        return HandleFirmwareUpload(id, params);
    };
}

nlohmann::json FirmwareUpgradeHandler::HandleFirmwareUpgrade(const std::string &id, const nlohmann::json &params)
{
    (void)params;
    // 更新状态
    currentStatus_.status = "upgrading";
    currentStatus_.progress = 0;
    currentStatus_.message = "正在升级固件...";
    currentStatus_.error_msg = "";

    auto &upgrader = Upgrade::IUpgradeService::GetInstance();
    // 完成流式升级
    if (!upgrader.CompleteUpgrade()) {
        currentStatus_.status = "error";
        currentStatus_.message = "升级失败: " + upgrader.GetLastErrorMessage();
        currentStatus_.error_msg = upgrader.GetLastErrorMessage();
        return CreateErrorResponse(id, 500, "升级完成失败: " + currentStatus_.error_msg);
    }

    // 返回升级完成状态，告知前端将要重启
    nlohmann::json result = {{"status", "success"}, {"message", "升级完成，系统将在10秒后重启"}};

    return CreateSuccessResponse(id, result);
}

nlohmann::json FirmwareUpgradeHandler::HandleGetUpgradeStatus(const std::string &id, const nlohmann::json &params)
{
    (void)params;
    EL_DEBUG("处理firmware.getStatus请求");

    // 构建符合协议的响应
    nlohmann::json result = {{"status", currentStatus_.status.empty() ? "idle" : currentStatus_.status},
                             {"progress", currentStatus_.progress},
                             {"message", currentStatus_.message}};

    return CreateSuccessResponse(id, result);
}

nlohmann::json FirmwareUpgradeHandler::HandleFirmwareUpload(const std::string &id, const nlohmann::json &params)
{
    // 检查必要参数
    if (!params.contains("chunk_index") || !params["chunk_index"].is_number() || !params.contains("total_chunks") ||
        !params["total_chunks"].is_number() || !params.contains("chunk_data") || !params["chunk_data"].is_string()) {
        return CreateErrorResponse(id, 400, "缺少必要参数或参数格式不正确");
    }

    int chunkIndex = params["chunk_index"];
    int totalChunks = params["total_chunks"];
    std::string chunkData = params["chunk_data"];

    // 第一个分片需要包含文件信息
    if (chunkIndex == 0) {
        if (!params.contains("file_info") || !params["file_info"].is_object()) {
            return CreateErrorResponse(id, 400, "缺少文件信息");
        }

        const auto &fileInfo = params["file_info"];
        if (!fileInfo.contains("name") || !fileInfo.contains("size") || !fileInfo.contains("version")) {
            return CreateErrorResponse(id, 400, "文件信息不完整");
        }

        // 初始化上传状态
        uploadStatus_.fileName = fileInfo["name"];
        uploadStatus_.totalSize = fileInfo["size"];
        uploadStatus_.totalChunks = totalChunks;
        uploadStatus_.version = fileInfo["version"];

        // 更新升级状态为下载中
        currentStatus_.status = "downloading";
        currentStatus_.progress = 0;
        currentStatus_.message = "正在下载固件文件...";

        // 启动升级模块
        auto &upgrader = Upgrade::IUpgradeService::GetInstance();
        if (!upgrader.InitializeUpgradeEnvironment()) {
            currentStatus_.status = "error";
            currentStatus_.message = "升级模块启动失败";
            return CreateErrorResponse(id, 500, "升级模块启动失败: " + upgrader.GetLastErrorMessage());
        }

        // 设置进度回调
        upgrader.RegisterProgressCallback([this](int progress, const std::string &status) {
            currentStatus_.progress = progress;
            currentStatus_.message = status.empty() ? 
                ("正在升级固件...进度: " + std::to_string(progress) + "%") : status;
        });
    }

    // 验证分片序号
    if (chunkIndex >= totalChunks) {
        return CreateErrorResponse(id, 400, "分片序号超出范围");
    }

    // 流式处理分片数据
    std::vector<uint8_t> binaryData = El::Utils::Base64Decode(chunkData);
    auto &upgrader = Upgrade::IUpgradeService::GetInstance();
    if (!upgrader.StreamUpgradeData(binaryData.data(), binaryData.size())) {
        currentStatus_.status = "error";
        currentStatus_.message = "流式处理分片数据失败";
        return CreateErrorResponse(id, 500, "流式处理分片数据失败: " + upgrader.GetLastErrorMessage());
    }
    // 计算下载进度 (chunkIndex从0开始，所以需要+1)
    int downloadProgress = static_cast<int>((static_cast<double>(chunkIndex + 1) / totalChunks) * 100);
    nlohmann::json result = {{"status", "success"}, {"progress", downloadProgress}};

    return CreateSuccessResponse(id, result);
}

} // namespace Protocol
} // namespace El